import os  
import shutil


class Recorder(object):
    """
    实验记录器类，负责管理实验过程中的文件保存、日志记录、代码备份等功能
    主要用于保存实验快照，确保实验可复现性
    """
    def __init__(self, snapshot_pref, ignore_folder):
        """
        初始化记录器
        :param snapshot_pref: 实验结果保存的根目录路径
        :param ignore_folder: 复制代码时需要忽略的文件夹名称
        """
        # 若保存目录不存在，则创建该目录
        if not os.path.isdir(snapshot_pref):
            os.mkdir(snapshot_pref)
        self.save_path = snapshot_pref  # 实验结果的根保存路径
        
        # 定义各类日志文件的路径
        self.log_file = self.save_path + "log.txt"  # 训练过程日志（如损失、精度等）
        self.readme = self.save_path + "README.md"  # 实验说明文档（可记录实验目的、结论等）
        self.opt_file = self.save_path + "opt.log"  # 命令行参数配置日志（保存训练时的所有参数）
        self.code_path = os.path.join(self.save_path, "code/")  # 代码快照保存路径（备份实验代码）

        # 若README文件已存在，则删除旧文件（避免覆盖）
        if os.path.isfile(self.readme):
            os.remove(self.readme)
        
        # 若代码快照目录不存在，则创建
        if not os.path.isdir(self.code_path):
            os.mkdir(self.code_path)
        
        # 复制当前项目代码到快照目录（忽略指定文件夹）
        self.copy_code(dst=self.code_path, ignore_folder=ignore_folder)

        # 打印保存路径，提示用户结果将存储的位置
        print("\n======> Result will be saved at: ", self.save_path)

    def copy_code(self, src="./", dst="./code/", ignore_folder='Exps'):
        """
        复制项目代码到指定目录（用于备份实验代码，确保可复现性）
        :param src: 源代码根目录（默认为当前工作目录）
        :param dst: 代码备份目标目录
        :param ignore_folder: 需要忽略的文件夹名称（如实验结果目录，避免递归复制）
        """
        import uuid  # 用于生成唯一标识符，避免目录重名
        
        # 若目标目录已存在，在目录名后添加UUID后缀（避免覆盖已有备份）
        if os.path.isdir(dst):
            dst = "/".join(dst.split("/")) + "code_" + str(uuid.uuid4()) + "/"
        
        # 收集所有需要复制的文件路径
        file_abs_list = []
        src_abs = os.path.abspath(src)  # 转换为绝对路径，确保跨平台兼容性
        
        # 遍历源目录下的所有文件和子目录
        for root, dirs, files in os.walk(src_abs):
            # 忽略包含指定文件夹的路径（如避免复制已保存的实验结果）
            if ignore_folder not in root:
                for name in files:
                    # 收集符合条件的文件绝对路径
                    file_abs_list.append(root + "/" + name)
        
        # 复制收集到的文件到目标目录
        for file_abs in file_abs_list:
            # 分割文件名和后缀，用于过滤不需要的文件
            file_split = file_abs.split("/")[-1].split('.')
            
            # 过滤条件：文件大小小于10MB（避免复制大文件）且不是.pyc文件（Python编译缓存文件）
            if os.path.getsize(file_abs)/1024/1024 < 10 and not file_split[-1] == "pyc":
                src_file = file_abs  # 源文件路径
                # 构建目标文件路径（保持与源目录相同的目录结构）
                dst_file = dst + file_abs.replace(src_abs, "")
                
                # 若目标文件的父目录不存在，则创建
                if not os.path.exists(os.path.dirname(dst_file)):
                    os.makedirs(os.path.dirname(dst_file))
                
                # 复制文件到目标路径
                try:
                    shutil.copyfile(src=src_file, dst=dst_file)
                except:
                    print("copy file error")  # 复制失败时打印错误提示

    def writeopt(self, opt):
        """
        保存命令行参数配置（如学习率、批次大小等）到opt.log
        :param opt: argparse解析得到的参数对象，包含所有训练配置
        """
        with open(self.opt_file, "w") as f:
            # 遍历参数对象的属性，写入文件
            for k, v in opt.__dict__.items():
                f.write(str(k)+": "+str(v)+"\n")

    def writelog(self, input_data):
        """
        写入训练过程日志（如每个epoch的损失、评估指标等）到log.txt
        :param input_data: 需要记录的日志内容（字符串或其他可转换为字符串的对象）
        """
        txt_file = open(self.log_file, 'a+')  # 以追加模式打开文件
        txt_file.write(str(input_data) + "\n")  # 写入内容并换行
        txt_file.close()

    def writereadme(self, input_data):
        """
        写入实验说明到README.md（如实验目的、关键结论、异常情况等）
        :param input_data: 需要记录的说明内容
        """
        txt_file = open(self.readme, 'a+')  # 以追加模式打开文件
        txt_file.write(str(input_data) + "\n")  # 写入内容并换行
        txt_file.close()

    # 以下方法为注释掉的预留功能（可能用于模型结构可视化或权重分析）
    def gennetwork(self, var):
        """生成模型网络结构（未实现，可能用于可视化）"""
        self.graph.draw(var=var)

    def savenetwork(self):
        """保存模型网络结构可视化图（未实现）"""
        self.graph.save(file_name=self.save_path+"network.svg")

    """def writeweights(self, input_data, block_id, layer_id, epoch_id):
        记录模型权重信息（如卷积层权重统计值）到日志文件（注释掉未使用）
        txt_path = self.weight_folder + "conv_weight_" + str(epoch_id) + ".log"
        txt_file = open(txt_path, 'a+')
        write_str = "%d\t%d\t%d\t" % (epoch_id, block_id, layer_id)
        for x in input_data:
            write_str += str(x) + "\t"
        txt_file.write(write_str+"\n")

    def drawhist(self):
        绘制权重分布直方图（注释掉未使用）
        drawer = DrawHistogram(txt_folder=self.weight_folder, fig_folder=self.weight_fig_folder)
        drawer.draw()"""
